#include <algorithm>
#include <cstdint>
#include <iostream>
#include <vector>
#include <sys/types.h>
#include <sys/stat.h>
#include <unistd.h>
#include <fstream>
#include <fcntl.h>

#include "acl/acl.h"
#include "aclnn_trace_custom.h"

#define SUCCESS 0
#define FAILED 1
#define INFO_LOG(fmt, args...) fprintf(stdout, "[INFO]  " fmt "\n", ##args)
#define ERROR_LOG(fmt, args...) fprintf(stderr, "[ERROR]  " fmt "\n", ##args)
#define CHECK_RET(cond, expr) do { if(!(cond)){ expr; } } while(0)

bool ReadFile(const std::string &filePath, void *buffer, size_t expectBytes)
{
    std::ifstream file(filePath, std::ios::binary);
    if (!file.is_open()) {
        ERROR_LOG("open %s failed", filePath.c_str());
        return false;
    }
    file.read(reinterpret_cast<char*>(buffer), expectBytes);
    if (file.gcount() != (std::streamsize)expectBytes) {
        ERROR_LOG("read size mismatch: expect=%zu got=%ld", expectBytes, (long)file.gcount());
        return false;
    }
    return true;
}

bool WriteFile(const std::string &filePath, const void *buffer, size_t size)
{
    std::ofstream ofs(filePath, std::ios::binary | std::ios::trunc);
    if(!ofs.is_open()){
        ERROR_LOG("open %s failed", filePath.c_str());
        return false;
    }
    ofs.write(reinterpret_cast<const char*>(buffer), size);
    return ofs.good();
}

int64_t GetShapeSize(const std::vector<int64_t> &shape)
{
    int64_t v = 1;
    for (auto d : shape) v *= d;
    return v;
}

int Init(int32_t deviceId, aclrtStream *stream)
{
    auto ret = aclInit(nullptr);
    CHECK_RET(ret == ACL_SUCCESS, return FAILED);
    ret = aclrtSetDevice(deviceId);
    CHECK_RET(ret == ACL_SUCCESS, return FAILED);
    ret = aclrtCreateStream(stream);
    CHECK_RET(ret == ACL_SUCCESS, return FAILED);
    return SUCCESS;
}

template <typename T>
int CreateAclTensor(const std::vector<int64_t> &shape, void **deviceAddr,
                    aclDataType dataType, aclTensor **tensor, const T *hostData = nullptr)
{
    size_t bytes = GetShapeSize(shape) * sizeof(T);
    auto ret = aclrtMalloc(deviceAddr, bytes, ACL_MEM_MALLOC_HUGE_FIRST);
    CHECK_RET(ret == ACL_SUCCESS, return FAILED);
    if (hostData != nullptr) {
        ret = aclrtMemcpy(*deviceAddr, bytes, hostData, bytes, ACL_MEMCPY_HOST_TO_DEVICE);
        CHECK_RET(ret == ACL_SUCCESS, return FAILED);
    }
    *tensor = aclCreateTensor(shape.data(), shape.size(), dataType,
                              nullptr, 0, ACL_FORMAT_ND,
                              shape.data(), shape.size(), *deviceAddr);
    return SUCCESS;
}

int main()
{
    int32_t deviceId = 0;
    aclrtStream stream;
    if (Init(deviceId, &stream) != SUCCESS) {
        ERROR_LOG("Init failed");
        return FAILED;
    }

    // 输入：形状 [batch, N, N]
    std::vector<int64_t> inputXShape = {1, 777, 777};
    int64_t inputElems = GetShapeSize(inputXShape);
    std::vector<aclFloat16> inputHost(inputElems);

    if (!ReadFile("../input/input_x.bin", inputHost.data(), inputElems * sizeof(aclFloat16))) {
        ERROR_LOG("read input failed");
        return FAILED;
    }

    // 输出：Trace 结果 => [batch]
    std::vector<int64_t> outputYShape = { inputXShape[0] };
    int64_t outputElems = GetShapeSize(outputYShape);
    std::vector<aclFloat16> outputHost(outputElems, (aclFloat16)0);

    void *inputDev = nullptr;
    void *outputDev = nullptr;
    aclTensor *inputTensor = nullptr;
    aclTensor *outputTensor = nullptr;

    CHECK_RET(CreateAclTensor(inputXShape, &inputDev, ACL_FLOAT16, &inputTensor, inputHost.data()) == SUCCESS, return FAILED);
    // 不传 hostData 参数，避免模板推导 nullptr 失败
    CHECK_RET((CreateAclTensor<aclFloat16>(outputYShape, &outputDev, ACL_FLOAT16, &outputTensor) == SUCCESS), return FAILED);
    // 输出用于 atomicAdd，显式清零
    auto memsetRet = aclrtMemset(outputDev, outputElems * sizeof(aclFloat16), 0x00, outputElems * sizeof(aclFloat16));
    CHECK_RET(memsetRet == ACL_SUCCESS, ERROR_LOG("memset output failed %d", memsetRet); return FAILED);
    // 获取 workspace（无 diagonal 属性）
    uint64_t workspaceSize = 0;
    aclOpExecutor *executor = nullptr;
    auto ret = aclnnTraceCustomGetWorkspaceSize(inputTensor, outputTensor, &workspaceSize, &executor);
    CHECK_RET(ret == ACL_SUCCESS, ERROR_LOG("GetWorkspaceSize failed %d", ret); return FAILED);

    void *workspaceDev = nullptr;
    if (workspaceSize > 0) {
        ret = aclrtMalloc(&workspaceDev, workspaceSize, ACL_MEM_MALLOC_HUGE_FIRST);
        CHECK_RET(ret == ACL_SUCCESS, ERROR_LOG("workspace alloc failed %d", ret); return FAILED);
    }

    ret = aclnnTraceCustom(workspaceDev, workspaceSize, executor, stream);
    CHECK_RET(ret == ACL_SUCCESS, ERROR_LOG("aclnnTraceCustom failed %d", ret); return FAILED);

    ret = aclrtSynchronizeStream(stream);
    CHECK_RET(ret == ACL_SUCCESS, ERROR_LOG("sync failed %d", ret); return FAILED);

    ret = aclrtMemcpy(outputHost.data(), outputElems * sizeof(aclFloat16),
                      outputDev, outputElems * sizeof(aclFloat16),
                      ACL_MEMCPY_DEVICE_TO_HOST);
    CHECK_RET(ret == ACL_SUCCESS, ERROR_LOG("memcpy D2H failed %d", ret); return FAILED);

    if (!WriteFile("../output/output_y.bin", outputHost.data(), outputElems * sizeof(aclFloat16))) {
        ERROR_LOG("write output failed");
    } else {
        INFO_LOG("trace result written: elements=%ld", outputElems);
    }

    aclDestroyTensor(inputTensor);
    aclDestroyTensor(outputTensor);
    if (inputDev) aclrtFree(inputDev);
    if (outputDev) aclrtFree(outputDev);
    if (workspaceDev) aclrtFree(workspaceDev);
    aclrtDestroyStream(stream);
    aclrtResetDevice(deviceId);
    aclFinalize();
    return SUCCESS;
}